import os, json
from jinja2 import Template
import jsonlines
from datetime import datetime
#from openai import AzureOpenAI
# from zhipuai import ZhipuAI
from openai import OpenAI
from typing import List, Dict, Any
from tqdm import tqdm

from loguru import logger
def json_loads(json_str: str, ensure_ascii: bool = False, use_json_repair: bool = True) -> Any:
    if use_json_repair:
        from json_repair import repair_json
        return repair_json(json_str, return_objects=True, ensure_ascii=ensure_ascii)
    else:
        try:
            return json.loads(json_str)
        except json.JSONDecodeError as e:
            logger.error(f"Error: {e}")
            return None

def json_load(json_fd , ensure_ascii: bool = False, use_json_repair: bool = True) -> Any:
    if use_json_repair:
        from json_repair import repair_json
        return repair_json(json_fd=json_fd, return_objects=True, ensure_ascii=ensure_ascii)
    else:
        try:
            return json.load(json_fd)
        except json.JSONDecodeError as e:
            logger.error(f"Error: {e}")
            return None

# 配置客户端

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_API_BASE_URL = os.getenv("OPENAI_API_BASE_URL")
DEPLOYMENT_NAME = os.getenv("OPENAI_MODEL_NAME")

class LLMClient:
    def __init__(self, api_key, base_url, deployment):
        # 使用 Azure OpenAI 客户端初始化
        # self.client = AzureOpenAI(
        #     azure_endpoint=base_url,
        #     api_key=api_key,
        #     api_version="2024-05-01-preview",
        # )
        self.client = OpenAI(
            base_url=base_url,
            api_key=api_key,
        )
        self.deployment = deployment

    def generate_dialogue(self, prompt):
        system_prompt = "你是一个对话生成器，请根据接收到的json数据，按照对话示例生成相应的多轮对话。"
        full_prompt = f"{system_prompt} {prompt}"

        # 调用 Azure OpenAI API 生成对话
        response = self.client.chat.completions.create(
            model=self.deployment,  # 请填写您要调用的模型名称
            messages=[{"role": "user", "content": full_prompt}],
            temperature=0.2,
            max_tokens=8192
        )
        return response.choices[0].message.content

# 读取 JSONL 文件并解析
def read_jsonl(file_path: str) -> List[Dict]:
    with open(file_path, "r", encoding="utf-8") as f:
        return [json_loads(line.strip()) for line in f]

# 动态生成对话内容
def generate_dialogue_content(api_definition: Dict, num_diags: int) -> str:
    prompt_dialogue = """下面是读取的jsonl数据：
{{ api_definition }}

下面是一组对话示例，对话过程中ASSISTANT使用到了给定的api:
<对话示例>
api定义:
{
    "name": "generate_password",
    "description": "Generate a random password",
    "parameters": {
        "type": "object",
        "properties": {
            "length": {
                "type": "integer",
                "description": "The length of the password"
            },
            "include_symbols": {
                "type": "boolean",
                "description": "Whether to include symbols in the password",
                "default": true
            }
        },
        "required": [
            "length"
        ]
    }
}
对话：
[
[
    {"from": "USER", "value": "I need a new password. Can you generate one for me?"},
    {"from": "ASSISTANT", "value": "Of course. How long would you like your password to be? And would you like it to include symbols?"},
    {"from": "USER", "value": " I would like it to be 12 characters long and yes, please include symbols."},
    {"from": "ASSISTANT", "value": "<functioncall> {\"name\": \"generate_password\", \"arguments\": {\"length\": 12, \"include_symbols\": true}}</functioncall>"},
    {"from": "FUNCTION RESPONSE", "value": "{\"password\": \"4&7j#9@1Q6*\"}"},
    {"from": "ASSISTANT", "value": "Here is your new password: 4&7j#9@1Q6*. Please make sure to save it in a secure location."}
]
]
<对话示例说明>
针对如上示例的几点说明：
（1）其中ASSISTANT回答中出现了<functioncall>, 代表要开始使用api了。<functioncall>和</functioncall>之间的信息为json的字符串表示，内容包含了要调用的api和传入参数："name"为ASSITANT要调用的api名字，"arguments"里包含了该api各传入参数的值
（2）FUNCTION RESPONSE的回答内容是api调用返回的结果
（3）FUNCTION RESPONSE之后的ASSISTANT的回答，是ASSISTANT理解FUNCTION RESPONSE之后的回答
</对话示例说明>

在理解了如上内容后，请根据如下要求生成内容。
要求：

1. api定义：
    {{ api_definition }}

2. 参数处理规则：
   (1) required参数不全时：
      ① 明确告知用户缺失的参数名称（格式：检测到缺少必填参数：xxx）
      ② 如果用户没有明确提供，不要自行猜测任何required参数的值
      ③ 得到所有必选参数值前，禁止调用functioncall

   (2) 默认参数处理：
      ① 如果用户未提供有default值的参数，直接使用默认值
      ② 需要明确告知用户默认值的具体内容

3. 对话质量要求：
   (1) FUNCTION RESPONSE内容必须具体，如：
      合格响应：{"status":"success", "order_id":"TS202308201234"}
      不合格响应：{"result":"操作成功"}

   (2) 用户不会直接提及api名称，需要通过上下文推理

4. 生成要求：
   生成{{ num_diags }}组对话，每组对话需满足：
   - 用户表达方式差异显著（如：有的用专业术语，有的用口语化表达）
   - 必须包含至少1个参数缺失需要猜测的对话场景
   - 必须包含1次用户直接提供所有参数的场景
   - JSON数组格式返回，不要其他描述
   - 严格保证JSON格式符合要求，字符串中的引号等必须加上转义符
   - 格式遵循<对话示例>中的对话部分，要包含USER, ASSISTANT的内容

5. 对话内容必须要包含FUNCTION RESPONSE的内容，并且这部分内容不能用占位符之类的含糊省略描述，可以想象生成合乎逻辑的api调用返回结果。
6. 对话内容连贯自然。对话过程中USER不能直接提到1中的api，是否调用api要靠ASSISTANT通过逻辑推理自行判断。
7. 若USER没有提供该api的required传入参数，ASSISTANT要主动询问USER这些参数的取值。如果有defalut取值的传入参数USER未提供，ASSISTANT要告知USER这些参数的dafault取值，否则不要结束对话。
8. 对话内容为中文。


"""
    t = Template(prompt_dialogue)
    prompt = t.render(api_definition=json.dumps(api_definition, ensure_ascii=False, indent=4), num_diags=num_diags)
    
    return prompt

def validate_dialogue(dialogue, api_definition):
    required_params = api_definition["parameters"].get("required", [])
    
    # 提取所有functioncall（带异常捕获）
    function_calls = []
    for turn in dialogue:
        try:
            # 检查是否包含必要字段
            if turn.get("from") == "ASSISTANT" and "<functioncall>" in turn.get("value", ""):
                function_calls.append(turn["value"])
        except (KeyError, AttributeError, TypeError) as e:
            # 跳过格式错误的对话轮次
            continue
    
    # 如果没有functioncall直接跳过
    if not function_calls:
        return
    
    # 检查最后一个functioncall（带多层异常捕获）
    try:
        last_call = function_calls[-1]
        
        # 提取参数部分
        args_str = last_call.split("<functioncall>")[1].split("</functioncall>")[0].strip()
        args = json.loads(args_str).get("arguments", {})
        
        # 检查必填参数
        missing = [p for p in required_params if p not in args]
        if missing:
            return  # 可记录日志或返回错误信息
            
        # 参数类型检查
        for param, spec in api_definition["parameters"]["properties"].items():
            if param in args:
                param_type = spec.get("type")
                value = args[param]
                
                if param_type == "integer" and not isinstance(value, int):
                    return
                elif param_type == "boolean" and not isinstance(value, bool):
                    return
                # 可扩展其他类型检查...
                
        return True  # 验证通过

    except (IndexError, KeyError, json.JSONDecodeError, AttributeError) as e:
        # 跳过所有格式错误的情况
        return

def get_args():
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument("--start", type=int, default=0)
    parser.add_argument("--end", type=int, default=100)
    parser.add_argument("--num_diags", type=int, default=2)
    return parser.parse_args()

# 主执行逻辑
def main():
    args = get_args()
    num_diags=args.num_diags
    api_data_list = read_jsonl("all_test_data.jsonl")
    #api_data_list = read_jsonl("data/round_a_test.jsonl")
    
    api_data_list = api_data_list[args.start:args.end]

  
    llm_client = LLMClient(OPENAI_API_KEY, OPENAI_API_BASE_URL, DEPLOYMENT_NAME)

    output_file_name = f"output_{args.start}_{args.end}.jsonl" 
    #output_file_name = "1.jsonl"
    with open(output_file_name, "w", encoding="utf-8") as output_file:
        for api_data in tqdm(api_data_list):
            # 遍历每个API
            for api in tqdm(api_data["apis"]):

                # 生成对话内容
                prompt = generate_dialogue_content(api, num_diags=num_diags)
                response = llm_client.generate_dialogue(prompt)
                #print(response)
                response = response.split("</think>")[-1].strip()
                response_content = response.replace("```json\n", "").replace("```", "")
                print(f"{response_content=}")

                try:
                    dialogues = json_loads(response_content)
                    if not isinstance(dialogues, list):
                        raise ValueError("生成结果不是列表格式")

                    for dialogue in dialogues:
                        try:
                            # 验证对话
                            validate_dialogue(dialogue, api)
                            
                            # 写入文件
                            output_line = json.dumps({
                                "api_name": api["name"],
                                "dialogue": dialogue
                            }, ensure_ascii=False)
                            
                            output_file.write(output_line + "\n")
                            output_file.flush()
                        
                        except ValueError as ve:
                            print(f"验证失败丢弃数据: {str(ve)}")
                            print("问题对话内容:", json.dumps(dialogue, indent=2, ensure_ascii=False))
                            
                except json.JSONDecodeError as e:
                    print(f"JSON解析失败: {str(e)}")
                    print("原始响应内容:", response_content)

    print("处理完成，结果已保存到 output.jsonl")

if __name__ == "__main__":
    main()
